import torch
import numpy as np
from pathlib import Path

'''
https://github.com/tancik/fourier-feature-networks/blob/master/Demo.ipynb
'''

def fourier_mapping(x, B):
    # a = torch.ones_like(B[:,0])
    # a = a/a.norm()
    # @ means regular matrix multiply
    x_proj = (2.*np.pi*x) @ B.transpose(0,1)
    return torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)

def generate_B_gauss(mapping_size, scale, use_cuda, B_file):
    if B_file.is_file():
        B_gauss=torch.load(B_file)
    else:
        B_gauss = torch.randn(mapping_size)*scale
        torch.save(B_gauss, B_file)
    if use_cuda:
        return B_gauss.cuda()
    return B_gauss

def get_ffn_embedder(mapping_size,scale,use_cuda,B_file):
    B_gauss = generate_B_gauss(mapping_size,scale,use_cuda,B_file)
    return lambda x:fourier_mapping(x, B_gauss), mapping_size[0]*2

if __name__ == '__main__':
    embed_fn, input_ch = get_ffn_embedder([64,2], 10.0, True, Path('./B_gauss.pth'))
    inputs_flat = torch.rand([8,128,2]).cuda()

    fourier_basis = embed_fn(inputs_flat)

    print('Done!')